# 计算一只股票所有日期的高低峰
import os

import numpy as np
import pandas as pd
from loguru import logger

from models.stock_model import StockNumber, DayInfo
from mylib.myredis import mr


def get_all_stock(stocks, N, end_days):
    logger.info('开始计算峰谷值')
    df = pd.read_csv('all.csv')
    for row in df.index:
        if stocks is not None and df.loc[row]['ts_code'] not in stocks:
            continue
        sn = StockNumber(df.loc[row])
        if '退' in sn.name:
            continue
        if 'ST' in sn.name:
            continue
        if not str(sn.ts_code).startswith('60') \
                and not str(sn.ts_code).startswith('30') \
                and not str(sn.ts_code).startswith('00'):
            continue
        analysis_stock(sn, N, end_days)
    logger.info('完成计算峰谷值')


def analysis_stock(sn, N, end_days):
    csv_path = f'stocks/{sn.ts_code}.csv'
    df2 = pd.read_csv(csv_path)
    q10 = []
    q11 = []
    all_di = []
    parr = 0
    today_di = None
    cnt = 0
    pct_arr = []
    arr_high = []
    arr_low = []
    date_arr = []
    price_arr = []
    delta_days = []
    delta_prices = []
    delta_d = 0
    price_arr_high = []
    price_arr_low = []

    for row2 in df2.index:
        di = DayInfo(sn, df2.loc[row2])
        price_arr_high.append(di.high)
        price_arr_low.append(di.low)
        d2 = mr.get_dict(sn.ts_code)
        if d2:
            if row2 > end_days:
                return
            while int(float(di.trade_date)) <= \
                    int(float(list(pd.DataFrame(d2).to_dict()['date'].values())[1])):
                del d2[0]
            d2_0 = d2[0]
            d2_1 = d2[1]
            price = di.low if float(d2_0['all_pct']) < 0 else di.high
            delta_price = round(price - float(d2_1['price']), 2)
            delta_day = int(d2_0['delta_days']) + 1
            pct = round(delta_price / float(d2_1['price']) * 100, 2)
            d2[0] = {
                'date': di.trade_date,
                'price': price,
                'delta_prices': delta_price,
                'delta_days': delta_day,
                'all_pct': pct,
            }
            cp = CalPoint(sn.ts_code, sn.name, **pd.DataFrame(d2).to_dict())
            cp.cal_data()
            continue
        delta_d += 1
        if len(q10) == N * 2 + 1:
            ad = all_di[N]
            if max(q10) == q10[N]:
                delta_days.append(-delta_d) if delta_days else delta_days.append(-(delta_d - 6))
                delta_d = 0
                delta_price = round(parr - q10[N], 2)
                delta_prices.append(delta_price)
                pct = round(delta_price / q10[N] * 100, 2)
                pct_arr.append(pct)
                if pct > 0:
                    arr_high.append(pct)
                else:
                    arr_low.append(pct)
                price_arr.append(q10[N])
                date_arr.append(ad.trade_date)
                parr = q10[N]
                cnt += 1
            elif min(q11) == q11[N]:
                delta_price = round(parr - q11[N], 2)
                delta_prices.append(delta_price)
                pct = round(delta_price / q11[N] * 100, 2)
                pct_arr.append(pct)
                if pct > 0:
                    delta_days.append(delta_d) if delta_days else delta_days.append(delta_d - 6)
                    arr_high.append(pct)
                else:
                    delta_days.append(-delta_d) if delta_days else delta_days.append(-(delta_d - 6))
                    arr_low.append(pct)
                delta_d = 0
                price_arr.append(q11[N])
                date_arr.append(ad.trade_date)
                parr = q11[N]
                cnt += 1
            all_di.pop(0)
            q10.pop(0)
            q11.pop(0)
        all_di.append(di)
        q10.append(di.high)
        q11.append(di.low)
        if today_di is None:
            parr = di.low
            today_di = di
            price_arr.append(di.low)
            date_arr.append(di.trade_date)
    data_dict = {
        'date': date_arr[:-1],
        'price': price_arr[:-1],
        'delta_prices': delta_prices,
        'delta_days': delta_days,
        'all_pct': pct_arr,
    }
    cp = CalPoint(sn.ts_code, sn.name, **data_dict)
    cp.cal_data()
    mr.set_dict(sn.ts_code, data_dict)


class CalPoint(object):

    def __init__(self, ts_code, name, date, price, delta_prices, delta_days, all_pct):
        self.ts_code = ts_code
        self.name = name
        self.date = list(date.values()) if isinstance(date, dict) else date
        self.price = list(price.values()) if isinstance(price, dict) else price
        self.delta_prices = list(delta_prices.values()) if isinstance(delta_prices, dict) else delta_prices
        self.delta_days = list(delta_days.values()) if isinstance(delta_days, dict) else delta_days
        self.all_pct = list(all_pct.values()) if isinstance(all_pct, dict) else all_pct

    def cal_data(self):
        # logger.info(self.ts_code)
        # logger.info(f'len=({len(self.date)})self.date \t\t\t= {self.date}')
        # logger.info(f'len=({len(self.price)})self.price \t\t= {self.price}')
        # logger.info(f'len=({len(self.delta_prices)})self.delta_prices \t= {self.delta_prices}')
        # logger.info(f'len=({len(self.delta_days)})self.delta_days \t= {self.delta_days}')
        # logger.info(f'len=({len(self.all_pct)})self.all_pct \t\t= {self.all_pct}')

        self.all_pct = [float(item) for item in self.all_pct]
        down_arr = [float(item) for item in self.all_pct if str(item).startswith('-')]
        up_arr = [float(item) for item in self.all_pct if not str(item).startswith('-')]

        median_high = round(np.median(up_arr), 2)
        median_low = round(np.median(down_arr), 2)

        # logger.info(f'all_pct avg_down\t= {median_high}')
        # logger.info(f'all_pct avg_up\t\t= {median_low}')
        msg = ''
        if self.all_pct[0] <= 0:
            if down_arr[0] <= median_low:
                msg = f'{self.date[0]},{self.ts_code},已低于平均跌幅,{median_low}%,跌,{down_arr[0]}%,{self.price[1]},{self.price[0]}'
                logger.success(msg)
            else:
                msg = f'{self.date[0]},{self.ts_code},未低于平均跌幅,{median_low}%,跌,{down_arr[0]}%,{self.price[1]},{self.price[0]}'
                logger.warning(msg)

        if self.all_pct[0] > 0:
            if up_arr[0] >= median_high:
                msg = f'{self.date[0]},{self.ts_code},已高于平均涨幅,{median_high}%,涨,{up_arr[0]}%,{self.price[1]},{self.price[0]}'
                logger.error(msg)
            else:
                msg = f'{self.date[0]},{self.ts_code},未高于平均涨幅,{median_high}%,涨,{up_arr[0]}%,{self.price[1]},{self.price[0]}'
                logger.warning(msg)

        csv_path = f"main_name_for2/{self.ts_code.replace('.', '_')}_{self.name}.csv"
        if not os.path.exists(csv_path):
            with open(csv_path, 'a+') as fw:
                # 20240119,603109.SH,未高于平均涨幅,14.73%,涨,2.27%,15.85,16.21
                fw.write('日期,code,des1,avg_pct,des2,now_pct,from_price,to_price')

        with open(f"main_name_for2/{self.ts_code.replace('.', '_')}_{self.name}.csv", 'a+') as fw:
            fw.write(msg)
            fw.write('\n')
